package org.zjx.message;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.zjx.message.messages.CommandMessage;
import org.zjx.message.messages.HeartbeatMessage;
import org.zjx.message.messages.RpcRequest;
import org.zjx.message.messages.RpcResponse;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;

public class KryoSerializer {
    private static final ThreadLocal<Kryo> kryoThreadLocal = ThreadLocal.withInitial(() -> {
        Kryo kryo = new Kryo();
        kryo.setRegistrationRequired(false);
        // 注册所有消息类型
        kryo.register(RpcRequest.class);
        kryo.register(RpcResponse.class);
        kryo.register(HeartbeatMessage.class);
        kryo.register(CommandMessage.class);
        // 其他消息类...

        // 配置字符串序列化器
        kryo.addDefaultSerializer(String.class, new StringSerializer());


        return kryo;
    });

    // 自定义字符串序列化器确保UTF-8编码
    public static class StringSerializer extends com.esotericsoftware.kryo.Serializer<String> {
        @Override
        public void write(Kryo kryo, Output output, String object) {
            output.writeString(object);
        }

        @Override
        public String read(Kryo kryo, Input input, Class<? extends String> aClass) {
            return input.readString();
        }

    }

    public static byte[] serialize(Object obj) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        Output output = new Output(baos);
        kryoThreadLocal.get().writeClassAndObject(output, obj);
        output.close();
        return baos.toByteArray();
    }

    @SuppressWarnings("unchecked")
    public static <T> T deserialize(byte[] bytes) {
        Input input = new Input(new ByteArrayInputStream(bytes));
        return (T) kryoThreadLocal.get().readClassAndObject(input);
    }
}